"""DHF模型参数率定示例

演示如何使用历史洪水数据对DHF模型进行参数率定。
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
from scipy.optimize import differential_evolution
from loguru import logger

from service.dhf_model import DHFModel
from service.dhf_model.dhf_utils import ModelUtils
from util.font_config import setup_chinese_font


class DHFCalibrator:
    """DHF模型参数率定器"""
    
    def __init__(self, model: DHFModel):
        self.model = model
        self.best_params = None
        self.best_performance = None
        self.calibration_history = []
    
    def objective_function(self, params, rainfall_data, observed_flow, evaporation_data=None):
        """
        目标函数：最大化NSE
        
        Args:
            params: 参数数组 [K, B, S0, U0, Ka, K0]
            rainfall_data: 降雨数据
            observed_flow: 观测流量
            evaporation_data: 蒸发数据
            
        Returns:
            负的NSE值 (因为优化算法是最小化)
        """
        try:
            # 更新模型参数
            K, B, S0, U0, Ka, K0 = params
            
            # 更新产流参数
            self.model.parameters.runoff_params['K'] = K
            self.model.parameters.runoff_params['B'] = B
            self.model.parameters.runoff_params['S0'] = S0
            self.model.parameters.runoff_params['U0'] = U0
            
            # 更新汇流参数
            self.model.parameters.routing_params['Ka'] = Ka
            self.model.parameters.routing_params['K0'] = K0
            
            # 重新初始化计算模块
            self.model.runoff_generation = self.model.runoff_generation.__class__(self.model.parameters)
            self.model.routing = self.model.routing.__class__(self.model.parameters, self.model.rainfall_center)
            self.model.routing.set_watershed_area(self.model.watershed_area)
            
            # 运行模拟
            results = self.model.run_simulation(rainfall_data, evaporation_data)
            
            # 计算性能指标
            simulated_flow = results['flow_series']
            min_len = min(len(observed_flow), len(simulated_flow))
            
            performance = ModelUtils.calculate_model_performance(
                np.array(observed_flow[:min_len]),
                simulated_flow[:min_len]
            )
            
            nse = performance['NSE']
            
            # 记录率定历史
            self.calibration_history.append({
                'params': params.copy(),
                'nse': nse,
                'performance': performance
            })
            
            # 返回负的NSE (因为要最小化)
            return -nse
            
        except Exception as e:
            logger.error(f"参数评估出错: {e}")
            return 1000  # 返回一个很大的值表示失败
    
    def calibrate_with_de(self, rainfall_data, observed_flow, evaporation_data=None,
                         max_iterations=100, population_size=15):
        """
        使用差分进化算法进行参数率定
        
        Args:
            rainfall_data: 降雨数据
            observed_flow: 观测流量
            evaporation_data: 蒸发数据
            max_iterations: 最大迭代次数
            population_size: 种群大小
            
        Returns:
            率定结果字典
        """
        logger.info("开始使用差分进化算法进行参数率定...")
        
        # 定义参数边界 [K, B, S0, U0, Ka, K0]
        bounds = [
            (0.1, 1.0),    # K: 下层下渗曲线曲率
            (0.5, 2.5),    # B: 下渗抛物线分配曲线形状参数
            (20, 120),     # S0: 表层蓄水容量
            (50, 150),     # U0: 下层蓄水容量
            (0.5, 0.95),   # Ka: 前期净雨影响程度衰减系数
            (0.2, 0.8)     # K0: 拟合线的斜率
        ]
        
        # 清空率定历史
        self.calibration_history = []
        
        # 运行差分进化算法
        result = differential_evolution(
            self.objective_function,
            bounds,
            args=(rainfall_data, observed_flow, evaporation_data),
            maxiter=max_iterations,
            popsize=population_size,
            seed=42,
            disp=True,
            callback=self._callback
        )
        
        # 保存最佳参数
        self.best_params = result.x
        best_nse = -result.fun
        
        logger.success(f"\n参数率定完成！")
        logger.info(f"最佳NSE: {best_nse:.4f}")
        logger.info("最佳参数:")
        param_names = ['K', 'B', 'S0', 'U0', 'Ka', 'K0']
        for i, (name, value) in enumerate(zip(param_names, self.best_params)):
            logger.info(f"  {name}: {value:.4f}")
        
        # 使用最佳参数重新运行模拟
        final_results = self.objective_function(
            self.best_params, rainfall_data, observed_flow, evaporation_data
        )
        
        return {
            'best_params': dict(zip(param_names, self.best_params)),
            'best_nse': best_nse,
            'optimization_result': result,
            'calibration_history': self.calibration_history
        }
    
    def _callback(self, xk, convergence):
        """优化过程回调函数"""
        if len(self.calibration_history) % 10 == 0:
            current_nse = -self.calibration_history[-1]['nse'] if self.calibration_history else 0
            logger.debug(f"迭代 {len(self.calibration_history)}: 当前最佳NSE = {current_nse:.4f}")
    
    # 第171行 - plot_calibration_history方法中
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    
    def plot_calibration_history(self, save_path=None):
        """
        绘制率定过程图
        """
        if not self.calibration_history:
            logger.warning("没有率定历史数据")
            return
        
        # 确保字体配置正确
        setup_chinese_font()
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
        
        # NSE变化过程
        nse_values = [record['nse'] for record in self.calibration_history]
        iterations = range(1, len(nse_values) + 1)
        
        ax1.plot(iterations, nse_values, 'b-', alpha=0.7, linewidth=1)
        ax1.set_xlabel('迭代次数')
        ax1.set_ylabel('NSE')
        ax1.set_title('参数率定过程 - NSE变化')
        ax1.grid(True, alpha=0.3)
        
        # 最佳NSE的累积最大值
        cumulative_best = np.maximum.accumulate(nse_values)
        ax1.plot(iterations, cumulative_best, 'r-', linewidth=2, label='累积最佳NSE')
        ax1.legend()
        
        # 参数变化过程 (显示前4个主要参数)
        param_names = ['K', 'B', 'S0', 'U0']
        colors = ['red', 'green', 'blue', 'orange']
        
        for i, (param_name, color) in enumerate(zip(param_names, colors)):
            param_values = [record['params'][i] for record in self.calibration_history]
            ax2.plot(iterations, param_values, color=color, alpha=0.7, 
                    linewidth=1, label=param_name)
        
        ax2.set_xlabel('迭代次数')
        ax2.set_ylabel('参数值')
        ax2.set_title('参数变化过程')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig


def create_multi_flood_data():
    """创建多场洪水的示例数据"""
    floods = []
    
    # 洪水1: 中等强度洪水
    flood1_rainfall = np.zeros(48)
    flood1_rainfall[8:20] = [2, 5, 12, 18, 25, 20, 15, 10, 8, 5, 3, 1]
    flood1_observed = create_flood_hydrograph(flood1_rainfall, peak_flow=400, base_flow=30)
    
    floods.append({
        'name': '洪水1 (中等强度)',
        'rainfall': flood1_rainfall.tolist(),
        'observed_flow': flood1_observed,
        'start_time': datetime(2023, 6, 15)
    })
    
    # 洪水2: 大洪水
    flood2_rainfall = np.zeros(72)
    flood2_rainfall[12:36] = [1, 3, 8, 15, 25, 35, 40, 38, 32, 28, 22, 18, 
                             15, 12, 10, 8, 6, 4, 3, 2, 1, 1, 0, 0]
    flood2_observed = create_flood_hydrograph(flood2_rainfall, peak_flow=1200, base_flow=45)
    
    floods.append({
        'name': '洪水2 (大洪水)',
        'rainfall': flood2_rainfall.tolist(),
        'observed_flow': flood2_observed,
        'start_time': datetime(2023, 7, 20)
    })
    
    # 洪水3: 小洪水
    flood3_rainfall = np.zeros(36)
    flood3_rainfall[6:18] = [1, 2, 5, 8, 12, 10, 8, 6, 4, 3, 2, 1]
    flood3_observed = create_flood_hydrograph(flood3_rainfall, peak_flow=200, base_flow=25)
    
    floods.append({
        'name': '洪水3 (小洪水)',
        'rainfall': flood3_rainfall.tolist(),
        'observed_flow': flood3_observed,
        'start_time': datetime(2023, 8, 10)
    })
    
    return floods


def create_flood_hydrograph(rainfall, peak_flow=500, base_flow=50):
    """根据降雨创建理想化的洪水过程线"""
    hours = len(rainfall) + 48  # 延长以包含完整退水过程
    flow = np.full(hours, base_flow)
    
    # 简单的降雨-径流转换
    total_rainfall = np.sum(rainfall)
    if total_rainfall > 0:
        # 找到降雨重心
        rainfall_center = np.sum(np.arange(len(rainfall)) * rainfall) / total_rainfall
        
        # 洪峰时间 (降雨重心后6-12小时)
        peak_time = int(rainfall_center + 8)
        
        # 构造洪水过程线
        for i in range(hours):
            if i < peak_time:
                # 涨水段
                if i > rainfall_center:
                    progress = (i - rainfall_center) / (peak_time - rainfall_center)
                    flow[i] = base_flow + (peak_flow - base_flow) * progress ** 1.5
            else:
                # 退水段
                recession_rate = 0.15  # 退水系数
                flow[i] = base_flow + (peak_flow - base_flow) * np.exp(-recession_rate * (i - peak_time))
    
    # 添加噪声
    noise = np.random.normal(0, flow * 0.03)
    flow = np.maximum(flow + noise, base_flow * 0.9)
    
    return flow.tolist()


def main():
    """主函数"""
    logger.info("DHF模型参数率定示例")
    logger.info("=" * 50)
    
    # 1. 创建模型实例
    logger.info("\n1. 创建DHF模型实例...")
    model = DHFModel(
        initial_flow=40.0,
        rainfall_center='中游',
        watershed_area=1500.0
    )
    
    # 2. 准备多场洪水数据
    logger.info("\n2. 准备多场洪水数据...")
    floods = create_multi_flood_data()
    
    for i, flood in enumerate(floods):
        logger.info(f"{flood['name']}: 总降雨量 {sum(flood['rainfall']):.1f}mm, "
              f"数据长度 {len(flood['rainfall'])}小时")
    
    # 3. 选择用于率定的洪水 (使用第一场洪水)
    logger.info("\n3. 使用第一场洪水进行参数率定...")
    calibration_flood = floods[0]
    
    # 4. 创建率定器并进行率定
    logger.info("\n4. 开始参数率定...")
    calibrator = DHFCalibrator(model)
    
    calibration_results = calibrator.calibrate_with_de(
        rainfall_data=calibration_flood['rainfall'],
        observed_flow=calibration_flood['observed_flow'],
        max_iterations=50,
        population_size=10
    )
    
    # 5. 显示率定结果
    logger.info("\n5. 率定结果:")
    logger.info(f"最佳NSE: {calibration_results['best_nse']:.4f}")
    logger.info("最佳参数:")
    for param, value in calibration_results['best_params'].items():
        logger.info(f"  {param}: {value:.4f}")
    
    # 6. 验证率定结果 (使用其他洪水)
    logger.info("\n6. 验证率定结果...")
    validation_results = []
    
    for i, flood in enumerate(floods[1:], 1):  # 跳过用于率定的第一场洪水
        logger.info(f"\n验证洪水 {i+1}: {flood['name']}")
        
        # 重置模型状态
        model.reset_model()
        
        # 运行模拟
        results = model.run_simulation(
            rainfall_data=flood['rainfall'],
            start_time=flood['start_time']
        )
        
        # 评估性能
        performance = model.evaluate_performance(flood['observed_flow'])
        validation_results.append({
            'flood_name': flood['name'],
            'performance': performance
        })
        
        logger.info(f"  NSE: {performance['NSE']:.4f}")
        logger.info(f"  R²: {performance['R2']:.4f}")
        logger.info(f"  RMSE: {performance['RMSE']:.2f}")
    
    # 7. 生成率定过程图
    logger.info("\n7. 生成率定过程图...")
    try:
        os.makedirs("./output", exist_ok=True)
        calibration_fig = calibrator.plot_calibration_history("./output/calibration_history.png")
        logger.success("率定过程图已保存到 ./output/calibration_history.png")
        plt.show()
    except Exception as e:
        logger.error(f"图表生成出现问题: {e}")
    
    # 8. 生成验证结果图表
    logger.info("\n8. 生成验证结果图表...")
    try:
        # 为每场验证洪水生成对比图
        for i, (flood, validation) in enumerate(zip(floods[1:], validation_results), 1):
            model.reset_model()
            results = model.run_simulation(
                rainfall_data=flood['rainfall'],
                start_time=flood['start_time']
            )
            model.evaluate_performance(flood['observed_flow'])
            
            plots = model.generate_plots(save_dir=f"./output/flood_{i+1}")
            logger.success(f"洪水{i+1}验证图表已保存到 ./output/flood_{i+1}/")
    
    except Exception as e:
        logger.error(f"验证图表生成出现问题: {e}")
    
    # 9. 生成率定报告
    logger.info("\n9. 生成率定报告...")
    try:
        report = generate_calibration_report(calibration_results, validation_results)
        with open("./output/calibration_report.txt", "w", encoding="utf-8") as f:
            f.write(report)
        logger.success("率定报告已保存到 ./output/calibration_report.txt")
    except Exception as e:
        logger.error(f"报告生成出现问题: {e}")
    
    logger.success("\n参数率定示例完成！")
    logger.info("\n建议:")
    logger.info("1. 使用更多历史洪水数据进行率定和验证")
    logger.info("2. 根据流域特性调整参数边界")
    logger.info("3. 考虑多目标优化 (如同时优化NSE、洪峰误差、峰现时间误差)")
    logger.info("4. 进行参数敏感性分析")


def generate_calibration_report(calibration_results, validation_results):
    """生成率定报告"""
    report = "DHF模型参数率定报告\n"
    report += "=" * 50 + "\n\n"
    
    # 率定结果
    report += "一、参数率定结果\n"
    report += "-" * 20 + "\n"
    report += f"率定NSE: {calibration_results['best_nse']:.4f}\n\n"
    
    report += "最佳参数:\n"
    for param, value in calibration_results['best_params'].items():
        report += f"  {param}: {value:.4f}\n"
    report += "\n"
    
    # 验证结果
    report += "二、模型验证结果\n"
    report += "-" * 20 + "\n"
    
    for validation in validation_results:
        report += f"{validation['flood_name']}:\n"
        perf = validation['performance']
        report += f"  NSE: {perf['NSE']:.4f}\n"
        report += f"  R²: {perf['R2']:.4f}\n"
        report += f"  RMSE: {perf['RMSE']:.2f}\n"
        report += f"  PBIAS: {perf['PBIAS']:.1f}%\n\n"
    
    # 总体评价
    avg_nse = np.mean([v['performance']['NSE'] for v in validation_results])
    report += "三、总体评价\n"
    report += "-" * 20 + "\n"
    report += f"验证期平均NSE: {avg_nse:.4f}\n"
    
    if avg_nse > 0.75:
        evaluation = "优秀"
    elif avg_nse > 0.65:
        evaluation = "良好"
    elif avg_nse > 0.50:
        evaluation = "满意"
    else:
        evaluation = "不满意"
    
    report += f"模型整体评价: {evaluation}\n"
    
    return report


if __name__ == "__main__":
    main()